/*
 * This file used to testing SM2/SM3/SM4 performance
 *
 * ===================================================
 *
 * author: linlh
 * date  : 2021.01.21
 * copyright: MIN-Group
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include "openssl/evp.h"
#include "openssl/bn.h"
#include "openssl/ec.h"
#include "openssl/sm2.h"
#include "openssl/sm3.h"

#define BENCHMARK_TIME 10000 // 测试次数

EC_KEY *generate_key() {
    EC_KEY *sm2key = NULL;

    // 这种方法就不需要在设置Group信息了，这个函数内会进行设置
    sm2key = EC_KEY_new_by_curve_name(OBJ_sn2nid("SM2"));
    /*
      OR
    sm2key = EC_KEY_new_by_curve_name(OBJ_sn2nid("sm2"));
      OR
    sm2key = EC_KEY_new_by_curve_name(NID_sm2);
    */
    if (!sm2key)
    {
        printf("Create SM2 Key Object error.\n");
        goto err;
    }
    
    if (EC_KEY_generate_key(sm2key) == 0)
    {
        printf("Error Of Generate SM2 Key.\n");
        goto err;
    }

    return sm2key;
    
err:
    if (sm2key) EC_KEY_free(sm2key);
}

EVP_PKEY *generate_key_evp() {
	const EC_GROUP *group = NULL;
    EVP_PKEY *sm2key = NULL;
    EVP_PKEY_CTX *pctx = NULL;
    unsigned char *out = NULL;
    size_t len;
    int loop;
    /*First Generate SM2 Key*/
    sm2key = EVP_PKEY_new();
    if (!sm2key)
    {   
        printf("Alloc EVP_PKEY Object error.\n");
        goto err;
    }   

    pctx = EVP_PKEY_CTX_new_id(EVP_PKEY_EC, NULL);
    if (!pctx)
    {   
        printf("Create EVP_PKEY_CTX Object error.\n");
        goto err;
    }   
    
    EVP_PKEY_keygen_init(pctx);
    if (!EVP_PKEY_CTX_set_ec_paramgen_curve_nid(pctx, NID_sm2))
    {   
        printf("Set EC curve name error.\n");
        goto err;
    }   
    
    if (!EVP_PKEY_CTX_set_ec_param_enc(pctx, OPENSSL_EC_NAMED_CURVE))
    {   
        printf("Set EC curve is named curve error.\n");
        goto err;
    }   
    
    if (EVP_PKEY_keygen(pctx, &sm2key) <= 0)
    {   
        printf("Generate SM2 key error.\n");
        goto err;
    }   

    /*OUTPUT EVP PKEY
    len = i2d_PrivateKey(sm2key, &out);
    if (len <= 0)
    {   
        printf("Output SM2 Private Key Error.\n");
        goto err;
    }   

    if (out) OPENSSL_free(out);
    out = NULL;

    len = i2d_PublicKey(sm2key, &out);
    if (len > 0)
    {
        printf("Generated SM2 Public Key ASN1 value: [");
        for (loop = 0; loop < len; loop++)
            printf("%02X", out[loop] & 0xff);
        printf("]\n");
    }

    if (out) OPENSSL_free(out);
    out = NULL;*/

    /*
    group = EC_KEY_get0_group(EVP_PKEY_get0_EC_KEY(sm2key));
    
    out = (unsigned char *)BN_bn2hex(EC_KEY_get0_private_key(EVP_PKEY_get0_EC_KEY(sm2key)));
    if (!out)
    {
        printf("Error Of Output SM2 Private key.\n");
        goto err;
    }

    //printf("\n              Private Key: [%s]\n", out);
    OPENSSL_free(out);
    out = (unsigned char *)EC_POINT_point2hex(group, EC_KEY_get0_public_key((EVP_PKEY_get0_EC_KEY(sm2key))), POINT_CONVERSION_UNCOMPRESSED, NULL);
    if (!out)
    {
        printf("Error Of Output SM2 Public key.\n");
        goto err;
    }
    //printf("              Public Key: [%s]\n", out);*/
    
     
err:
    if (pctx) EVP_PKEY_CTX_free(pctx);
    if (out) OPENSSL_free(out);
	return sm2key;
}

void benchmark_generate_key() {
    const EC_GROUP *sm2group = NULL;
    EC_KEY *sm2key = NULL;
    char *out = NULL;
    size_t len;
    int loop;
    
    sm2key = EC_KEY_new_by_curve_name(OBJ_sn2nid("SM2"));
    /*
      OR
    sm2key = EC_KEY_new_by_curve_name(OBJ_sn2nid("sm2"));
      OR
    sm2key = EC_KEY_new_by_curve_name(NID_sm2);
    */
    if (!sm2key)
    {
        printf("Create SM2 Key Object error.\n");
        goto err;
    }
    
    if (EC_KEY_generate_key(sm2key) == 0)
    {
        printf("Error Of Generate SM2 Key.\n");
        goto err;
    }
    
    sm2group = EC_KEY_get0_group(sm2key);
    
    /*Output SM2 Key*/
    out = BN_bn2hex(EC_KEY_get0_private_key(sm2key));
    if (!out)
    {
        printf("Error Of Output SM2 Private key.\n");
        goto err;
    }

    // printf("Generated SM2 Private Key: [%s]\n", out);
    OPENSSL_free(out);
    out = EC_POINT_point2hex(sm2group, EC_KEY_get0_public_key(sm2key), POINT_CONVERSION_UNCOMPRESSED, NULL);
    if (!out)
    {
        printf("Error Of Output SM2 Public key.\n");
        goto err;
    }
    // printf("              Public Key: [%s]\n", out);
err:
    if (sm2key) EC_KEY_free(sm2key);
    if (out) OPENSSL_free(out);
}

void benchmark_generate_key_evp() {
    const EC_GROUP *group = NULL;
	EVP_PKEY *sm2key = NULL;
	EVP_PKEY_CTX *pctx = NULL;
	unsigned char *out = NULL;
	size_t len;
	int loop;
	/*First Generate SM2 Key*/
	sm2key = EVP_PKEY_new();
	if (!sm2key)
	{
		printf("Alloc EVP_PKEY Object error.\n");
		goto err;
	}

	pctx = EVP_PKEY_CTX_new_id(EVP_PKEY_EC, NULL);
	if (!pctx)
	{
		printf("Create EVP_PKEY_CTX Object error.\n");
		goto err;
	}
    
	EVP_PKEY_keygen_init(pctx);
	if (!EVP_PKEY_CTX_set_ec_paramgen_curve_nid(pctx, NID_sm2))
	{
		printf("Set EC curve name error.\n");
		goto err;
	}
    
	if (!EVP_PKEY_CTX_set_ec_param_enc(pctx, OPENSSL_EC_NAMED_CURVE))
	{
		printf("Set EC curve is named curve error.\n");
		goto err;
	}
    
	if (EVP_PKEY_keygen(pctx, &sm2key) <= 0)
	{
		printf("Generate SM2 key error.\n");
		goto err;
	}

    /*OUTPUT EVP PKEY*/
    len = i2d_PrivateKey(sm2key, &out);
    if (len <= 0)
    {
        printf("Output SM2 Private Key Error.\n");
        goto err;
    }
    
    /*printf("Generated SM2 Private Key ASN1 value: [");
    for (loop = 0; loop < len; loop++)
        printf("%02X", out[loop] & 0xff);
    printf("]\n");*/

    if (out) OPENSSL_free(out);
    out = NULL;
       
    len = i2d_PublicKey(sm2key, &out);
    /*if (len > 0)
    {
        printf("Generated SM2 Public Key ASN1 value: [");
        for (loop = 0; loop < len; loop++)
            printf("%02X", out[loop] & 0xff);
        printf("]\n");
    }*/
    
    if (out) OPENSSL_free(out);
    out = NULL;

    /*OUTPUT X + Y + d*/
    group = EC_KEY_get0_group(EVP_PKEY_get0_EC_KEY(sm2key));
    
    /*Output SM2 Key*/
    out = (unsigned char *)BN_bn2hex(EC_KEY_get0_private_key(EVP_PKEY_get0_EC_KEY(sm2key)));
    if (!out)
    {
        printf("Error Of Output SM2 Private key.\n");
        goto err;
    }

    //printf("\n              Private Key: [%s]\n", out);
    OPENSSL_free(out);
    out = (unsigned char *)EC_POINT_point2hex(group, EC_KEY_get0_public_key((EVP_PKEY_get0_EC_KEY(sm2key))), POINT_CONVERSION_UNCOMPRESSED, NULL);
    if (!out)
    {
        printf("Error Of Output SM2 Public key.\n");
        goto err;
    }
    //printf("              Public Key: [%s]\n", out);

err:
	if (sm2key) EVP_PKEY_free(sm2key);
	if (pctx) EVP_PKEY_CTX_free(pctx);
	if (out) OPENSSL_free(out);
}


void benchmark_sign() {
    // 生成公私钥
    EC_KEY *sm2key = NULL;
    unsigned char *sig = NULL;
    unsigned char digest[SM3_DIGEST_LENGTH];
    size_t dgst_len;
    int loop, siglen;
     
    char userID[20] = "1234567812345678";
    char msg[8] = "MIN-TEAM";

    sm2key = generate_key();
    if (sm2key == NULL) {
        printf("sm2 key generate failed\n");
        goto err;
    }

    clock_t start, finish;
	double totaltime;

    start = clock();
    ECDSA_SIG *sig_der = NULL;
    // 连续签名
    for (int i = 0; i < BENCHMARK_TIME; ++i) {
	    sig_der = sm2_do_sign(sm2key, EVP_sm3(), userID,
    	                      strlen(userID), msg, 8);
    }
    finish = clock();
	totaltime = (double)(finish - start) / CLOCKS_PER_SEC / BENCHMARK_TIME;
	printf("签名时间测试（测试次数：%d）: %fs\n",BENCHMARK_TIME, totaltime);

    // 这里进行验签，输入的结果还是sm2key，为什么不是公钥？原因是函数内部会取用公钥的信息
    // 没有用到私钥的信息，这个函数实现的不好
    if (sm2_do_verify(sm2key, EVP_sm3(), sig_der, userID,
                      strlen(userID), msg,
                      8) != 1) {
        printf("sm2_do_verify fail!\n");
        goto err;
    }
err:
	OPENSSL_clear_free(sig, siglen);
    if (sm2key) EC_KEY_free(sm2key);
}

void benchmark_verify() {
    // 生成公私钥
    EC_KEY *sm2key = NULL;
    unsigned char *sig = NULL;
    unsigned char digest[SM3_DIGEST_LENGTH];
    size_t dgst_len;
    int loop, siglen;
     
    char userID[20] = "1234567812345678";
    char msg[8] = "MIN-TEAM";

    sm2key = generate_key();
    if (sm2key == NULL) {
        printf("sm2 key generate failed\n");
        goto err;
    }

    clock_t start, finish;
	double totaltime;

    ECDSA_SIG *sig_der = NULL;
	sig_der = sm2_do_sign(sm2key, EVP_sm3(), userID,
    	                      strlen(userID), msg, 8);

    start = clock();
    for (int i = 0; i < BENCHMARK_TIME; ++i) {
        if (sm2_do_verify(sm2key, EVP_sm3(), sig_der, userID,
                          strlen(userID), msg,
                          8) != 1) {
            printf("sm2_do_verify fail!\n");
            goto err;
        }
    }
    finish = clock();
	totaltime = (double)(finish - start) / CLOCKS_PER_SEC / BENCHMARK_TIME;
	printf("验签时间测试（测试次数：%d）: %fs\n",BENCHMARK_TIME, totaltime);
err:
	OPENSSL_clear_free(sig, siglen);
    if (sm2key) EC_KEY_free(sm2key);

}

// binary to string
int b2s(char *bin, char *outs) {
    int i = 0;
    char tmpbuf[4];
    int iRet = 0;
    char *ptr = bin;
    for (i = 0; i < strlen(bin) / 2; i++) {
        memset(tmpbuf, 0x00, sizeof(tmpbuf));
        memcpy(tmpbuf, ptr, 2);
        ptr += 2;
        iRet = strtol(tmpbuf, NULL, 16);
#ifndef NO_DEBUG
// printf("the iRet =[%d]\n", iRet);
#endif

        memset(outs++, iRet, 1);
    }
    return i;
}

EC_KEY *CalcSm2PublicKey(const char *pubkey_hex_string, char *private_hex_x) {
    int bn_len = 0;
    char raw_buf[128] = {0};
    BIGNUM *k = NULL;

    EC_KEY *ec_key = NULL;
    EC_POINT *pubkey = NULL;
    EC_GROUP *ec_group = NULL;

    ec_group = EC_GROUP_new_by_curve_name(OBJ_sn2nid("SM2"));
    if (ec_group == NULL) goto err;

    pubkey = EC_POINT_hex2point(ec_group, (const char *)pubkey_hex_string, NULL,
                                NULL);
    if (!pubkey) goto err;

    ec_key = EC_KEY_new();
    if (!ec_key) goto err;
    if (!EC_KEY_set_group(ec_key, ec_group)) {
        EC_KEY_free(ec_key);
        ec_key = NULL;
        goto err;
    }

    if (!EC_KEY_set_public_key(ec_key, pubkey)) {
        EC_KEY_free(ec_key);
        ec_key = NULL;
        goto err;
    }

    if (private_hex_x != NULL) {
        bn_len = b2s((char *)private_hex_x, raw_buf);
        printf("bn_len = [%d]\n", bn_len);

        k = BN_new();
        if (BN_bin2bn((const unsigned char *)raw_buf, bn_len, k) == NULL) {
            printf("bin2bn fail!\n");
            exit(0);
        }

        if (!EC_KEY_set_private_key(ec_key, k)) {
            EC_KEY_free(ec_key);
            ec_key = NULL;
            goto err;
        }
    }

err:
    if (pubkey) EC_POINT_free(pubkey);
    if (ec_group) EC_GROUP_free(ec_group);

    return ec_key;
}

EC_KEY *CalculatePubKey(const EC_GROUP *ec_group, const char *pub_hex_string) {
    EC_KEY *ec_key = NULL;
    EC_POINT *pubkey = NULL;

    if ((pubkey = EC_POINT_new(ec_group)) == NULL) goto err;
    if (!EC_POINT_hex2point(ec_group, pub_hex_string, pubkey, NULL)) goto err;

    if (!ec_key) {
        ec_key = EC_KEY_new();
        if (!ec_key) goto err;
        if (!EC_KEY_set_group(ec_key, ec_group)) {
            EC_KEY_free(ec_key);
            ec_key = NULL;
            goto err;
        }
    }

    if (!EC_KEY_set_public_key(ec_key, pubkey)) {
        EC_KEY_free(ec_key);
        ec_key = NULL;
        goto err;
    }

err:
    if (pubkey) EC_POINT_free(pubkey);

    return ec_key;
}

EC_KEY *CalculateKey(const EC_GROUP *ec_group, const char *privkey_hex_string) {
    EC_KEY *ec_key = NULL;
    EC_POINT *pubkey = NULL;
    BIGNUM *privkey = NULL;

    if (!BN_hex2bn(&privkey, (const char *)privkey_hex_string)) return NULL;
    if ((pubkey = EC_POINT_new(ec_group)) == NULL) goto err;
    if (!ec_key) {
        ec_key = EC_KEY_new();
        if (!ec_key) goto err;
        if (!EC_KEY_set_group(ec_key, ec_group)) {
            EC_KEY_free(ec_key);
            ec_key = NULL;
            goto err;
        }
    }

    if (!EC_POINT_mul(ec_group, pubkey, privkey, NULL, NULL, NULL)) {
        EC_KEY_free(ec_key);
        ec_key = NULL;
        goto err;
    }

    if (!EC_KEY_set_private_key(ec_key, privkey) ||
        !EC_KEY_set_public_key(ec_key, pubkey)) {
        EC_KEY_free(ec_key);
        ec_key = NULL;
        goto err;
    }

err:
    if (privkey) BN_free(privkey);
    if (pubkey) EC_POINT_free(pubkey);

    return ec_key;
}


#define DEBUG_CHAR_HEX(buf_ptr, buf_len)                   \
    printf("\n%s=[", #buf_ptr);                            \
    for (int i = 0; i < (buf_len); i++) {                      \
        printf("%02X", *((unsigned char *)(buf_ptr) + i)); \
    }                                                      \
    printf("]\n");

char *randstr(char *str, const int len) {
    srand(time(NULL));
    int i;
    for (i = 0; i < len; ++i)
    {
        switch ((rand() % 3))
        {
        case 1:
            str[i] = 'A' + rand() % 26;
            break;
        case 2:
            str[i] = 'a' + rand() % 26;
            break;
        default:
            str[i] = '0' + rand() % 10;
            break;
        }
    }
    str[++i] = '\0';
    return str;
}

void benchmark_encrypt_decrypt(int size) {
    EC_KEY *sm2key = NULL;
    const EC_GROUP *sm2group = NULL;
    char *privkey = NULL;
    char *pubkey = NULL;
    int retval;
    char *msg = (char *)malloc(size); // 用于存储测试用的消息

    int cryptsize = 0;
    if (size < 1024)
        cryptsize = 1024;
    else if (size < 8000)
        cryptsize = 4096;
    else 
        cryptsize = 20000;
    char *ciphertext_buf = (char *)malloc(cryptsize);  // 用于存储加密后的结果
    memset(ciphertext_buf, 0, cryptsize);
    unsigned char *ptext_buf = (unsigned char *)malloc(cryptsize);
    memset(ptext_buf, 0, cryptsize);
    unsigned char *in = NULL;

    randstr(msg, size-1);
    
    size_t ciphertext_len = 0;
    size_t ptext_len = 0;
    size_t inlen = 0; 

    sm2key = generate_key();
    if (sm2key == NULL) {
        printf("sm2 key generate failed!\n");
        goto err;
    }

    // 这里生成的sm2key并不是16进制的字符串，需要进行转换
    // 注意这里我们为什么要进行转换，因为实际应用中，我们读取的是公私钥的字符串信息
    // 而不是计算过程的ECC曲线的点或者其他信息
    // 公私钥也是要保存成字符串的
    privkey = BN_bn2hex(EC_KEY_get0_private_key(sm2key));
    if (!privkey) {
        printf("Error of translate SM2 Private key.\n");
        goto err;
    }
    //printf("SM2 private key: [%s]\n", privkey);

    sm2group = EC_KEY_get0_group(sm2key);

    pubkey = EC_POINT_point2hex(sm2group, EC_KEY_get0_public_key(sm2key), POINT_CONVERSION_UNCOMPRESSED, NULL);
    if (!pubkey) {
        printf("Error of translate SM2 Public key.\n");
        goto err;
    }
    //printf("SM2 public key : [%s]\n", pubkey);

    // 用公钥加密
    sm2key = CalcSm2PublicKey(pubkey, NULL); // 把公钥字符串转换成EC_KEY
    if (!sm2key) {
        printf("Error of Calculate SM2 Public Key.\n");
        goto err;
    }

    clock_t start, finish;
    double totaltime;
    start = clock();
    for (int i = 0; i < BENCHMARK_TIME; ++i) { 
        retval = sm2_encrypt(sm2key, EVP_sm3(), msg, strlen(msg), ciphertext_buf, (size_t *)&ciphertext_len);
        if (!retval) {
            printf("Error of calculate cipher text length.\n");
            goto err;
        }
    }
    finish = clock();
    totaltime = (double)(finish - start) / CLOCKS_PER_SEC / BENCHMARK_TIME;
    printf("SM2加密%d Bytes时间（测试次数：%d）: %fs\n", size-1, BENCHMARK_TIME, totaltime);
 

/////////////////////////////////////////////////////////////////////////////////////////

    // 将加密结果转化成16进制字符串
    char *encrypt_buf = (char *)malloc(cryptsize);
    memset(encrypt_buf, 0, cryptsize);

    for (int i = 0; i < ciphertext_len; ++i) {
        sprintf(encrypt_buf + i*2, "%02X", *((unsigned char *)(ciphertext_buf) + i));
    }

    //printf("SM2 Encrypt : [%s]\n", encrypt_buf);
////////////////////////////////////////////////////////////////////////////////////////

    // 私钥解密
    sm2group = EC_GROUP_new_by_curve_name(NID_sm2);
    if (!sm2group) {
        goto err;
    }

    sm2key = CalculateKey((const EC_GROUP *)sm2group, privkey);
    if (!sm2key) {
        printf("Error of Calculate SM2 Private Key.\n");
        goto err;
    }

    inlen = strlen(encrypt_buf) / 2;
    in = OPENSSL_malloc(inlen);
    if (!in) {
        printf("Error Of Alloc Memory.\n");
        goto err;
    }
    
    ciphertext_len = b2s(encrypt_buf, in);
    
    start = clock();
    for (int i = 0; i < BENCHMARK_TIME; ++i) {
        retval = sm2_decrypt(sm2key, EVP_sm3(), in, ciphertext_len, ptext_buf,
                                 &ptext_len);
        if (!retval) {
            printf("Error Of sm2_decrypt [%d].\n", retval);
            goto err;
        }
    }
    finish = clock();
    totaltime = (double)(finish - start) / CLOCKS_PER_SEC / BENCHMARK_TIME;
    printf("SM2解密%d Bytes时间（测试次数：%d）: %fs\n",size-1, BENCHMARK_TIME, totaltime);

    //printf("Decrypt: [%s]\n", ptext_buf); 
    if (strcmp(msg, ptext_buf) != 0) {
        printf("msg:[%s]. ptext_buf:[%s]\n", msg, ptext_buf);
        printf("Error of sm2 encrypt and decrypt don't match!\n");
        goto err;
    }

err:
    if (sm2key) EC_KEY_free(sm2key);
    if (privkey) OPENSSL_free(privkey);
    if (pubkey) OPENSSL_free(pubkey);
    if (in) OPENSSL_free(in);
    if (msg) free(msg);
    if (ciphertext_buf) free(ciphertext_buf);
    if (ptext_buf) free(ptext_buf);
    if (encrypt_buf) free(encrypt_buf);
}

int main(int argc, char *argv[]) {
    EC_KEY *sm2key = NULL;
	EVP_PKEY *sm2key_evp = NULL;

    clock_t start, finish;
	double totaltime;
#if 1
    // 公私钥生成性能验证
    start = clock();
    // 公私钥生成的时间较长，所以这里只测试10000万次
    for (int i = 0; i < BENCHMARK_TIME/10; ++i) {
        //benchmark_generate_key();
        sm2key = generate_key();
        if (sm2key == NULL) {
            printf("generate sm2 key error!\n");
            exit(0);
        } else {
            EC_KEY_free(sm2key);
        }
    }
	finish = clock();
	totaltime = (double)(finish - start) / CLOCKS_PER_SEC / BENCHMARK_TIME * 10;
	printf("公私钥生成时间（测试次数：%d）: %fs\n",BENCHMARK_TIME/10, totaltime);
#endif 

#if 1
    // 公私秘钥生成性能验证EVP接口
	start = clock();
    // 公私钥生成的时间较长，所以这里只测试10000万次
    for (int i = 0; i < BENCHMARK_TIME/10; ++i) {
        //benchmark_generate_key_evp();
        sm2key_evp = generate_key_evp();
        if (sm2key_evp == NULL) {
            printf("generate sm2 key error!\n");
            exit(0);
        } else {
            EVP_PKEY_free(sm2key_evp);
        }
    }
	finish = clock();
	totaltime = (double)(finish - start) / CLOCKS_PER_SEC / BENCHMARK_TIME * 10;
	printf("公私钥EVP接口生成时间（测试次数：%d）: %fs\n",BENCHMARK_TIME/10, totaltime);
#endif

#if 1 
    // 签名性能验证
    benchmark_sign();
#endif

#if 1 
    // 签名性能验证
    benchmark_verify();
#endif
    
    // 加解密性能验证
    benchmark_encrypt_decrypt(9);
    benchmark_encrypt_decrypt(1025);
    benchmark_encrypt_decrypt(8193);

	return 0;
}
